# -*- coding: utf-8 -*-

"""
LSTM难度预测模型训练器
"""

import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from tensorflow.keras.utils import to_categorical
import sqlite3
import os


class LSTMDifficultyTrainer:
    def __init__(self, db_path="data/game.db"):
        """初始化训练器"""
        self.db_path = db_path
        self.model = None
        
    def load_training_data(self):
        """
        从数据库加载训练数据
        
        Returns:
            tuple: (features, labels) 训练数据和标签
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 查询玩家游戏记录
        cursor.execute('''
            SELECT player_id, time_spent, moves_count, difficulty, completed
            FROM game_records
            ORDER BY player_id, created_at
        ''')
        
        records = cursor.fetchall()
        conn.close()
        
        # 组织数据：为每个玩家构建7天的历史记录
        player_data = {}
        for record in records:
            player_id, time_spent, moves_count, difficulty, completed = record
            if player_id not in player_data:
                player_data[player_id] = []
            
            # 将难度转换为数值标签
            difficulty_map = {"简单": 0, "中等": 1, "困难": 2, "专家": 3}
            label = difficulty_map.get(difficulty, 1)
            
            player_data[player_id].append({
                'time_spent': time_spent,
                'moves_count': moves_count,
                'completed': int(completed),
                'label': label
            })
        
        # 构建训练数据
        features = []
        labels = []
        
        for player_id, data in player_data.items():
            # 为每个玩家构建序列数据
            for i in range(len(data) - 7):
                # 取7天的历史数据作为特征
                sequence = data[i:i+7]
                feature_seq = [[item['time_spent'], item['moves_count'], item['completed']] 
                              for item in sequence]
                
                # 使用第8天的难度作为标签
                label = data[i+7]['label']
                
                features.append(feature_seq)
                labels.append(label)
        
        if len(features) == 0:
            # 如果没有足够的数据，生成一些示例数据用于演示
            print("警告：数据库中没有足够的训练数据，生成示例数据")
            features = np.random.rand(100, 7, 3) * [1800, 100, 1]  # 时间、步数、完成状态
            labels = np.random.randint(0, 4, 100)  # 4个难度等级
        else:
            features = np.array(features)
            labels = np.array(labels)
        
        # 对标签进行one-hot编码
        labels = to_categorical(labels, num_classes=4)
        
        return features, labels
    
    def build_model(self):
        """构建LSTM模型"""
        model = Sequential([
            LSTM(64, return_sequences=True, input_shape=(7, 3)),
            Dropout(0.2),
            LSTM(32, return_sequences=False),
            Dropout(0.2),
            Dense(32, activation='relu'),
            Dense(4, activation='softmax')  # 4个难度等级
        ])
        
        model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy']
        )
        
        self.model = model
        return model
    
    def train(self, epochs=50, validation_split=0.2):
        """
        训练模型
        
        Args:
            epochs (int): 训练轮数
            validation_split (float): 验证集比例
        """
        # 加载数据
        features, labels = self.load_training_data()
        
        # 构建模型
        if self.model is None:
            self.build_model()
        
        # 训练模型
        history = self.model.fit(
            features, labels,
            epochs=epochs,
            validation_split=validation_split,
            batch_size=32,
            verbose=1
        )
        
        return history
    
    def save_model(self, model_path="data/difficulty_model.tflite"):
        """
        保存模型为TensorFlow Lite格式
        
        Args:
            model_path (str): 模型保存路径
        """
        # 确保目录存在
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        
        # 转换为TensorFlow Lite模型
        converter = tf.lite.TFLiteConverter.from_keras_model(self.model)
        tflite_model = converter.convert()
        
        # 保存模型
        with open(model_path, 'wb') as f:
            f.write(tflite_model)
        
        print(f"模型已保存到: {model_path}")
    
    def evaluate_model(self):
        """评估模型性能"""
        features, labels = self.load_training_data()
        
        if self.model is None:
            print("模型未训练")
            return
        
        # 评估模型
        loss, accuracy = self.model.evaluate(features, labels, verbose=0)
        print(f"模型评估 - 损失: {loss:.4f}, 准确率: {accuracy:.4f}")
        
        return loss, accuracy


if __name__ == "__main__":
    # 使用示例
    trainer = LSTMDifficultyTrainer()
    
    # 构建并训练模型
    trainer.build_model()
    history = trainer.train(epochs=20)
    
    # 评估模型
    trainer.evaluate_model()
    
    # 保存为TensorFlow Lite格式
    trainer.save_model("data/difficulty_model.tflite")